{
 "cells": [
  {
   "cell_type": "code",
   "id": "3fea87d3-23fc-40b6-a38a-52d535322aaf",
   "metadata": {},
   "source": [
    "import cartopy.crs as ccrs\n",
    "import cartopy.feature as cfeature\n",
    "import matplotlib.pyplot as plt\n",
    "import xarray as xr"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "11ba477c-1e4a-48ab-9a5d-0e1b2a2433c2",
   "metadata": {},
   "source": [
    "# Load the NetCDF files"
   ]
  },
  {
   "cell_type": "code",
   "id": "4601bd05-9ad4-4613-a878-ae3b85fd42fa",
   "metadata": {},
   "source": [
    "ds_precip = xr.open_dataset(\"/Users/james.a/Downloads/chirps_monthly_cleaned_06.nc\")\n",
    "ds_spi_gammma_03 = xr.open_dataset(\"/Users/james.a/git/climate_indices/trial_chirps7_spi_gamma_03.nc\")\n",
    "ds_spi_pearson_03 = xr.open_dataset(\"/Users/james.a/git/climate_indices/trial_chirps7_spi_pearson_03.nc\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "d8c8c169-9051-44bf-9c07-0214686c022c",
   "metadata": {},
   "source": [
    "# Inspect the dataset"
   ]
  },
  {
   "cell_type": "code",
   "id": "a599f4b7-fc96-4e1b-922b-bca796e6be7e",
   "metadata": {},
   "source": [
    "# Display variables, dimensions, coordinates\n",
    "print(ds_precip)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "5fa22b97-d81f-4718-8560-ba5ea829c9d7",
   "metadata": {},
   "source": [
    "# select a time step by index\n",
    "step = 6\n",
    "time_step_precip = ds_precip[\"precip\"].isel(time=step)\n",
    "time_step_spi_gamma_03 = ds_spi_gammma_03[\"spi_gamma_03\"].isel(time=step)\n",
    "time_step_spi_pearson_03 = ds_spi_pearson_03[\"spi_pearson_03\"].isel(time=step)\n",
    "\n",
    "# create a map plot\n",
    "fig = plt.figure(figsize=(10, 6))  # Set figure size\n",
    "ax = plt.axes(projection=ccrs.PlateCarree())  # Use PlateCarree for simple lat/lon projection (common for grids)\n",
    "\n",
    "time_step = time_step_spi_pearson_03\n",
    "\n",
    "# Plot the 2D grid as a filled contour (or use .pcolormesh() for pixelated view)\n",
    "time_step.plot.contourf(\n",
    "    ax=ax,  # Attach to cartopy axis\n",
    "    transform=ccrs.PlateCarree(),  # Ensure data is projected correctly\n",
    "    cmap=\"viridis\",  # Colormap (options: 'coolwarm', 'jet', 'RdBu_r', etc.)\n",
    "    levels=20,  # Number of color levels (adjust for detail)\n",
    "    add_colorbar=True,  # Automatically add colorbar\n",
    "    cbar_kwargs={\"label\": \"Temperature (°C)\"},  # Customize colorbar\n",
    ")\n",
    "\n",
    "# Step 4: Add map features\n",
    "ax.coastlines()  # Add coastlines\n",
    "ax.add_feature(cfeature.BORDERS, linestyle=\":\")  # Add country borders\n",
    "ax.add_feature(cfeature.LAND, edgecolor=\"black\")  # Add land fill\n",
    "ax.add_feature(cfeature.OCEAN)  # Add ocean fill\n",
    "ax.gridlines(draw_labels=True)  # Add lat/lon gridlines with labels\n",
    "\n",
    "# Customize the plot\n",
    "plt.title(\"Temperature Map at Time: \" + str(time_step.time.values))  # Dynamic title with time stamp\n",
    "plt.show()"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "2d242ed3-73b6-4579-9b2a-90dc79a72c48",
   "metadata": {},
   "source": [
    "from IPython.display import HTML\n",
    "from matplotlib.animation import FuncAnimation\n",
    "\n",
    "# Select the variable to animate (e.g., spi_pearson_03; you can change this to another one)\n",
    "data_var = ds_spi_pearson_03[\"spi_pearson_03\"]  # Full DataArray with all time steps\n",
    "num_steps = data_var.sizes[\"time\"]  # Number of time steps\n",
    "\n",
    "# Create the figure and axis with cartopy projection\n",
    "fig = plt.figure(figsize=(10, 6))\n",
    "ax = plt.axes(projection=ccrs.PlateCarree())\n",
    "\n",
    "# Add static map features (these don't change per frame)\n",
    "ax.coastlines()\n",
    "ax.add_feature(cfeature.BORDERS, linestyle=\":\")\n",
    "ax.add_feature(cfeature.LAND, edgecolor=\"black\")\n",
    "ax.add_feature(cfeature.OCEAN)\n",
    "ax.gridlines(draw_labels=True)\n",
    "\n",
    "# Initialize the plot with the first time step using pcolormesh (easier to update than contourf)\n",
    "time_step_init = data_var.isel(time=0)\n",
    "im = time_step_init.plot.pcolormesh(\n",
    "    ax=ax,\n",
    "    transform=ccrs.PlateCarree(),\n",
    "    cmap=\"viridis\",  # Colormap\n",
    "    add_colorbar=False,  # We'll add a static colorbar below\n",
    ")\n",
    "\n",
    "# Add a static colorbar (based on the full data range for consistency)\n",
    "vmin = data_var.min().values  # Min value across all time steps\n",
    "vmax = data_var.max().values  # Max value across all time steps\n",
    "cbar = plt.colorbar(im, ax=ax, label=\"SPI Pearson 03\", extend=\"both\")  # Adjust label as needed\n",
    "im.set_clim(vmin, vmax)  # Set consistent color limits\n",
    "\n",
    "# Set initial title\n",
    "title = ax.set_title(\"Map at Time: \" + str(time_step_init.time.values))\n",
    "\n",
    "\n",
    "# Define the update function for each frame\n",
    "def update(frame):\n",
    "    # Select the current time step\n",
    "    time_step = data_var.isel(time=frame)\n",
    "\n",
    "    # Update the image data (pcolormesh array)\n",
    "    im.set_array(time_step.values.ravel())  # Flatten the 2D array for pcolormesh update\n",
    "\n",
    "    # Update the title with the current time\n",
    "    title.set_text(\"Map at Time: \" + str(time_step.time.values))\n",
    "\n",
    "    return [im, title]\n",
    "\n",
    "\n",
    "# Create the animation\n",
    "anim = FuncAnimation(\n",
    "    fig,\n",
    "    update,\n",
    "    frames=range(num_steps),  # Loop over all time steps\n",
    "    interval=500,  # 500 ms = 0.5 seconds per frame\n",
    "    repeat=True,  # Loop indefinitely\n",
    "    blit=True,  # Optimize by only redrawing changed elements\n",
    ")\n",
    "\n",
    "# Display the animation in Jupyter notebook\n",
    "HTML(anim.to_jshtml())"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "710df657-6b04-4704-b27b-8f1783647dd6",
   "metadata": {},
   "source": [],
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
